"""
根据在银行实际项目需求生成可供实验的数据
"""

import pandas as pd
import numpy as np
import os
import yaml
from utils import read_yaml, get_file_name
from torch_geometric.datasets import Planetoid
import torch

config = read_yaml(file_path="./config.yaml")["data_config"]


# device = config["model_config"]["device"]


def output_experiment_data(num_users, save_path="./dataset.pkl"):
    max_seq_len = config["event"]["max_seq_len"]
    event_vocab_size = config["event"]["vocab_size"]
    action_vocab_size = config["action"]["vocab_size"]
    max_view_times = config["action"]["max_view_times"]
    df = pd.DataFrame(columns=['id', 'event', 'action', 'label', 'flag', 'features', 'page_features'])
    df.id = range(num_users)
    df.event = [np.random.randint(1, event_vocab_size, size=(np.random.randint(1, max_seq_len),)).tolist() for _ in
                range(num_users)]
    df.action = [np.random.randint(1 + event_vocab_size, action_vocab_size + event_vocab_size,
                                   size=(np.random.randint(max_view_times),)).tolist()
                 for _ in
                 range(num_users)]
    df.label = np.random.randint(0, 2, size=(num_users, 1))
    df.sample(frac=1).reset_index(drop=True, inplace=True)
    df.flag = ["Train" for i in range(int(config["train_rate"] * num_users))] + ["Valid" for i in range(
        int(config["valid_rate"] * num_users))] + ["Test" for _ in range(
        num_users - int(config["valid_rate"] * num_users) - int(config["train_rate"] * num_users))]
    # df.features = [np.random.normal(size=(config["num_features"],)) for _ in range(num_users)]
    df.features = np.random.normal(size=(num_users, config["num_features"],)).tolist()
    df.page_features = np.random.normal(size=(num_users, config["num_features"],)).tolist()
    df.to_pickle(save_path)
    print("Generate Bank Data for Event and Action Successfully !")


output_experiment_data(num_users=2000)

# print(get_file_name())
